From 5615310d2139374bd7fdac0169babff33de529d8 Mon Sep 17 00:00:00 2001 From: "Karl O. Pinc" Date: Fri, 20 Sep 2024 22:27:12 -0500 Subject: [PATCH] Make downloading CSV work --- src/pgwui_sql/views/sql.py | 73 ++++++++++++++++++++++++++++++++------ tests/views/test_sql.py | 11 ++++-- 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/src/pgwui_sql/views/sql.py b/src/pgwui_sql/views/sql.py index 7e0d398..c169cea 100644 --- a/src/pgwui_sql/views/sql.py +++ b/src/pgwui_sql/views/sql.py @@ -19,9 +19,14 @@ from pyramid.view import view_config import attrs +import codecs +import csv +import datetime import logging import markupsafe import psycopg.errors +import pyramid.response +import tempfile import pgwui_core.core import pgwui_core.utils @@ -29,6 +34,8 @@ import pgwui_sql.views.base from pgwui_common.view import auth_base_view from pgwui_sql import exceptions as sql_ex +from pgwui_core.constants import CSV +from pgwui_sql.constants import MANY_FILES log = logging.getLogger(__name__) @@ -79,6 +86,7 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler): sql_results = attrs.field(factory=list) search_path = attrs.field(default=None) # requested search_path db_search_path = attrs.field(default=None) # search_path of db + tfile = attrs.field(default=None) def make_form(self): return pgwui_sql.views.base.SQLForm().build( @@ -89,6 +97,14 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler): super().read() self.search_path = self.request.POST.get('search_path') + def dl_filename(self): + uf = self.uf + return '_'.join( + ['sql_results', + uf["db"], + uf["user"], + datetime.datetime.now().isoformat(sep="_", timespec="seconds")]) + def write(self, result, errors): ''' Setup dict to render resulting html form @@ -105,6 +121,8 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler): response['filename'] = self.uf['filename'] response['report_success'] = (not response['errors'] and self.uf['action'] != '') + if self.uf['download']: + response['dl_filename'] = self.dl_filename() return response def val_input(self): @@ -181,11 +199,25 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler): return sql return f'SET search_path TO {self.search_path};\n{sql}' + def make_download(self, sql_results): + downloading = self.uf['download'] + if downloading: + self.tfile = tempfile.TemporaryFile(mode='w+t', newline='') + writer = csv.writer(self.tfile) + for sql_result in sql_results: + writer.writerow(sql_result.statusmessage.data) + writer.writerow(sql_result.heading.data) + for row in sql_result.rows: + writer.writerow(row.data) + writer.writerow(f'{sql_result.rowcount.data} rows') + self.tfile.seek(0) + def cleanup(self): ''' Execute a series of SQL statements. - The result goes into the upload handler (uh.sql_results), - interleaving errors with output. + The result goes into the upload handler (uh.sql_results). + This happens even when a file is built to download so that + should there be errors the results that do exist are displayed. ''' cur = self.cur null_rep = self.uf['null_rep'] @@ -213,6 +245,8 @@ class SQLResultsHandler(pgwui_core.core.SessionDBHandler): sql_results.append(sql_result) nextset = cur.nextset() + self.make_download(sql_results) + def factory(self, ue): '''Make a db loader function from an UploadEngine. @@ -234,14 +268,14 @@ def log_response(response): else: download_fmt = 'TAB' if pgwui_core.utils.is_checked(response['one_file_checked']): - download_as = 'One file' + download_as = 'one file' else: - download_as = 'A zip file of one file per result set' + download_as = 'a zip file of one file per result set' log.info(': '.join( [f'Successful sql: DB {response["db"]}', # ' File ({filename}):' f'Download Format {download_fmt}', - f'Download AS {download_as}', + f'Download as {download_as}', f'By user {response["user"]}'])) @@ -253,11 +287,28 @@ def sql_view(request): uh = SQLResultsHandler(request).init() response = pgwui_core.core.UploadEngine(uh).run() - settings = request.registry.settings - response.setdefault('pgwui', dict()) - response['pgwui']['pgwui_sql'] = settings['pgwui']['pgwui_sql'] + if uh.uf['download'] and not response['errors']: + pmd_response = pyramid.response.Response() + pmd_response.cache_control = 'public; max-age=0' + if uh.uf['download_as'] == MANY_FILES: + pmd_response.content_type = 'application/zip' + elif uh.uf['download_fmt'] == CSV: + pmd_response.content_type = 'text/csv' + else: + pmd_response.content_type = 'text/plain' + pmd_response.content_disposition = \ + f'attachment; filename={response["dl_filename"]}.csv' + + pmd_response.app_iter = codecs.iterencode(uh.tfile, 'utf_8') - response['result_rows'] = uh.sql_results + log_response(response) + return pmd_response + else: + settings = request.registry.settings + response.setdefault('pgwui', dict()) + response['pgwui']['pgwui_sql'] = settings['pgwui']['pgwui_sql'] - log_response(response) - return response + response['result_rows'] = uh.sql_results + + log_response(response) + return response diff --git a/tests/views/test_sql.py b/tests/views/test_sql.py index 00625de..f54b85b 100644 --- a/tests/views/test_sql.py +++ b/tests/views/test_sql.py @@ -53,9 +53,13 @@ CHANGED_RESPONSE = { 'trim_upload': constants.CHECKED, 'upload_null': constants.CHECKED, 'user': 'someuser', + 'download': False, } -UNCHANGED_RESPONSE = {'report_success': False} +UNCHANGED_RESPONSE = { + 'report_success': False, + 'download': False, +} # Helper classes @@ -78,11 +82,14 @@ def isolate_sql_view(monkeypatch, pyramid_request_config): cause UploadEngine to return the supplied "response". ''' def run(response): - def upload_engine(*args): + def upload_engine(uh, *args): + uh.read() return MockUploadEngine(response) monkeypatch.setattr( sql.pgwui_core.core, 'UploadEngine', upload_engine) + monkeypatch.setattr( + sql, 'log_response', lambda *args: None) settings = pyramid_request_config.get_settings() settings['pgwui'] = settings.get('pgwui', dict()) -- 2.34.1